package com.lnn.algorithm.aco;

import com.lnn.pojo.Node;
import com.lnn.pojo.Route;
import com.lnn.util.Address;
import com.lnn.util.Matrix;
import lombok.Data;
import org.springframework.web.multipart.MultipartFile;

import java.io.*;
import java.util.ArrayList;
import java.util.Objects;

/**
 * 蚁群算法实现类
 *
 * @author: Li Ning Ning
 * @time: 2021/7/23 16:21
 * @description: TODO
 * @version: 1.0
 */
@Data
public class ACO {

    /**
     * 城市数量
     */
    private int cityNum;

    /**
     * 蚂蚁数量
     */
    private int antNum;

    /**
     * 蚁群
     */
    private Ant[] ants;

    /**
     * 迭代次数
     */
    private int iterMax;

    /**
     * 信息素重要程度因子
     */
    private double alpha;

    /**
     * 启发函数重要程度因子
     */
    private double beta;

    /**
     * 信息素挥发因子
     */
    private double rho;

    /**
     * 信息素常数，表示蚂蚁遍历一次所有城市所释放的信息素总量
     */
    private double Q;

    /**
     * 信息素矩阵
     */
    private double[][] pheromone;

    /**
     * 距离矩阵
     */
    private int[][] distance;

    /**
     * 最佳路径编号
     */
    private int[] bestRoute;

    /**
     * 最佳长度
     */
    private int bestRouteLength;

    /**
     * 最佳路径地址
     */
    private String[] address;

    /**
     * 城市坐标x
     */
    private int[] x;

    /**
     * 城市坐标y
     */
    private int[] y;

    public ACO() {

    }

    public ACO(int antNum, int iterMax, double alpha, double beta, double rho, double Q) {
        this.antNum = antNum;
        this.ants = new Ant[antNum];
        this.iterMax = iterMax;
        this.alpha = alpha;
        this.beta = beta;
        this.rho = rho;
        this.Q = Q;
    }

    public ACO(int cityNum, int antNum, int iterMax, double alpha, double beta, double rho, double Q) {
        this.cityNum = cityNum;
        this.antNum = antNum;
        this.ants = new Ant[antNum];
        this.iterMax = iterMax;
        this.alpha = alpha;
        this.beta = beta;
        this.rho = rho;
        this.Q = Q;
        distance = new int[cityNum][cityNum];
    }

    /**
     * 初始化参数
     *
     * @param distance 距离矩阵
     * @param address  位置数组
     */
    public ACO init(int[][] distance, String[] address) {
        //初始化距离矩阵
        this.distance = distance;
        //初始化地址
        this.address = address;
        //初始化参数
        initParams();
        return this;
    }

    /**
     * 文件初始化参数
     */
    public ACO init(MultipartFile multipartFile) throws IOException {
        // 读取数据
        File file = new File(Objects.requireNonNull(multipartFile.getOriginalFilename()));
        readFile(multipartFile, file);

        String strBuff;
        BufferedReader data = new BufferedReader(new InputStreamReader(
                new FileInputStream(file)));
        //初始化城市总数
        String temp = data.readLine();
        this.cityNum = Integer.parseInt(temp);
        distance = new int[cityNum][cityNum];
        x = new int[cityNum];
        y = new int[cityNum];

        //初始化城市坐标
        for (int i = 0; i < cityNum; i++) {
            strBuff = data.readLine();
            String[] strCol = strBuff.split(" ");
            x[i] = Integer.parseInt(strCol[1]);
            y[i] = Integer.parseInt(strCol[2]);
        }
        // 计算距离矩阵
        this.distance = Matrix.computeDistanceMatrix(x, y);
        //初始化地址
        address = new String[cityNum];
        for (int i = 0; i < cityNum; i++) {
            address[i] = "Node" + i;
        }
        //初始化参数
        initParams();
        return this;
    }

    private void readFile(MultipartFile multipartFile, File file) {
        int n;
        try (InputStream in = multipartFile.getInputStream(); OutputStream os = new FileOutputStream(file)) {
            // 得到文件流。以文件流的方式输出到新文件
            byte[] buffer = new byte[4096];
            while ((n = in.read(buffer, 0, 4096)) != -1) {
                os.write(buffer, 0, n);//写入文件
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 初始化参数
     */
    private void initParams() {
        //初始化最优路径
        bestRoute = new int[cityNum + 1];
        //初始化最优路径长度
        bestRouteLength = Integer.MAX_VALUE;
        // 初始化信息素矩阵
        initPheromone();
        //随机放置蚂蚁
        initAnt();
    }

    /**
     * 初始化信息素矩阵
     */
    private void initPheromone() {
        pheromone = new double[cityNum][cityNum];
        for (int i = 0; i < cityNum; i++) {
            for (int j = 0; j < cityNum; j++) {
                pheromone[i][j] = 0.1; // 初始化为0.1
            }
        }
    }

    /**
     * 初始化随机放置蚂蚁
     */
    private void initAnt() {
        for (int i = 0; i < antNum; i++) {
            ants[i] = new Ant(cityNum);
            ants[i].init(distance, alpha, beta);
        }
    }

    /**
     * 多次迭代路径
     *
     * @return 路径列表
     */
    public ArrayList<Route> solve() {
        for (int iter = 0; iter < iterMax; iter++) {
            // 每一只蚂蚁移动的过程
            for (int i = 0; i < antNum; i++) {
                for (int j = 1; j < cityNum; j++) {
                    ants[i].selectNextCity(pheromone);
                }
                // 蚂蚁回到起始位置FirstCity
                ants[i].getTabu().add(ants[i].getFirstCity());
                // 计算蚂蚁路径的长度
                if (ants[i].getTourLength() < bestRouteLength) {
                    bestRouteLength = ants[i].getTourLength();
                    for (int k = 0; k < cityNum + 1; k++) {
                        bestRoute[k] = ants[i].getTabu().get(k);
                    }
                }
                for (int j = 0; j < cityNum; j++) {
                    ants[i].getDelta()[ants[i].getTabu().get(j)][ants[i]
                            .getTabu().get(j + 1)] = Q / ants[i]
                            .getTourLength();
                    ants[i].getDelta()[ants[i].getTabu().get(j + 1)][ants[i]
                            .getTabu().get(j)] = Q / ants[i]
                            .getTourLength();
                }
            }
            // 更新信息素
            updatePheromone();
            // 重新初始化蚂蚁
            for (int i = 0; i < antNum; i++) {
                ants[i].init(distance, alpha, beta);
            }
        }
        //格式化数组使得从初始点开始排序
        ArrayList<Route> routes = formatData();
        // 打印最佳结果
        printOptimal();
        return routes;
    }

    /**
     * 更新信息素
     */
    private void updatePheromone() {
        // 信息素挥发
        for (int i = 0; i < cityNum; i++) {
            for (int j = 0; j < cityNum; j++) {
                pheromone[i][j] = pheromone[i][j] * (1 - rho);
            }
        }
        // 信息素更新
        for (int i = 0; i < cityNum; i++) {
            for (int j = 0; j < cityNum; j++) {
                for (int k = 0; k < antNum; k++) {
                    pheromone[i][j] += ants[k].getDelta()[i][j];
                }
            }
        }
    }

    private ArrayList<Route> formatData() {
        bestRoute = Address.StartAtZero(bestRoute);
        address = Address.arrayToAddress(bestRoute, address);
        ArrayList<Route> routes = new ArrayList<>();
        for (int i = 0; i < bestRoute.length; i++) {
            Route route = new Route();
            route.setNumber(bestRoute[i]);
            route.setAddress(address[i]);
            route.setDistance(distance[bestRoute[i]][bestRoute[(i + 1) % bestRoute.length]]);
            route.setTotalCity(address.length - 1);
            route.setTotalDistance(bestRouteLength);
            if (x != null && y != null) {
                route.setNode(new Node(x[bestRoute[i]], y[bestRoute[i]]));
            }
            routes.add(route);
        }
        return routes;
    }

    private void printOptimal() {
        System.out.println("蚁群算法: ");
        System.out.println("The optimal length is: " + bestRouteLength);
        System.out.println("The optimal tour is: ");
        for (int i = 0; i < cityNum + 1; i++) {
            System.out.print(bestRoute[i] + " ");
        }
        System.out.println();
    }

}
